{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Step1: Create OpenVINO Runtime Core\n",
    "from openvino.runtime import Core\n",
    "core = Core()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Step2: Compile the Model, CompiledModel allow you to get information inputs\n",
    "# or output ports by a tensor name or index\n",
    "net = core.compile_model(\"yolov5s.xml\",\"AUTO\",{\"PERFORMANCE_HINT\": \"LATENCY\"})\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<ConstOutput: names[images] shape{1,3,640,640} type: f32>\n",
      "{1, 3, 640, 640} <Type: 'float32'> 0\n",
      "{'images'} <class 'set'> images <class 'str'>\n"
     ]
    }
   ],
   "source": [
    "# Learn the information of input node\n",
    "input_node = net.inputs[0]\n",
    "print(input_node)\n",
    "print(input_node.shape, input_node.element_type, input_node.index)\n",
    "print(input_node.names, type(input_node.names), input_node.any_name, type(input_node.any_name))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<ConstOutput: names[output] shape{1,25200,85} type: f32>\n"
     ]
    }
   ],
   "source": [
    "# Learn the information of output node\n",
    "output_node = net.outputs[0]\n",
    "print(output_node)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(720, 1280, 3)\n",
      "(640, 640, 3) (0.5, 0.5) 0.0 140.0\n"
     ]
    }
   ],
   "source": [
    "# Prepare the input data\n",
    "import cv2\n",
    "from utils.augmentations import letterbox\n",
    "img = cv2.imread(\"data/images/zidane.jpg\")\n",
    "#img = cv2.imread(\"data/images/bus.jpg\")\n",
    "print(img.shape)\n",
    "letterbox_img, ratio, (dw, dh) = letterbox(img, auto=False)\n",
    "print(letterbox_img.shape, ratio, dw, dh)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1, 3, 640, 640)"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Normalization + Swap RB + Layout from HWC to NCHW\n",
    "blob = cv2.dnn.blobFromImage(letterbox_img, 1/255.0, swapRB=True)\n",
    "blob.shape\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(1, 25200, 85)\n"
     ]
    }
   ],
   "source": [
    "# Step 3-5: Do the inference\n",
    "outs = net({input_node:blob})[output_node]\n",
    "print(outs.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {
    "tags": [
     "outputPrepend",
     "outputPrepend",
     "outputPrepend",
     "outputPrepend",
     "outputPrepend",
     "outputPrepend",
     "outputPrepend",
     "outputPrepend",
     "outputPrepend"
    ]
   },
   "outputs": [],
   "source": [
    "# Step6: Process inference results\n",
    "import numpy as np\n",
    "class_ids = []\n",
    "confidences = []\n",
    "boxes = []\n",
    "\n",
    "for out in outs:\n",
    "    for i, det in enumerate(out):\n",
    "        confidence = det[4]\n",
    "        scores = det[5:]\n",
    "        class_id = np.argmax(scores)\n",
    "        if scores[class_id] > 0.25:\n",
    "            confidences.append(confidence)\n",
    "            class_ids.append(class_id)\n",
    "            x,y,w,h = det[0].item(),det[1].item(),det[2].item(),det[3].item()\n",
    "            left = int((x - 0.5*w -dw) / ratio[0])\n",
    "            top = int((y - 0.5*h - dh) / ratio[1])\n",
    "            width = int(w / ratio[0])\n",
    "            height = int(h / ratio[1])\n",
    "            box = np.array([left, top, width, height])\n",
    "            boxes.append(box)\n",
    "\n",
    "indexes = cv2.dnn.NMSBoxes(boxes, confidences, 0.25, 0.45)\n",
    "\n",
    "filtered_ids = []\n",
    "filered_confidences = []\n",
    "filtered_boxes = []\n",
    "\n",
    "for i in indexes:\n",
    "    filtered_ids.append(class_ids[i])\n",
    "    filered_confidences.append(confidences[i])\n",
    "    filtered_boxes.append(boxes[i])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [],
   "source": [
    "import yaml\n",
    "with open(\"data/coco.yaml\",\"r\",encoding=\"utf-8\") as f:\n",
    "    result = yaml.load(f.read(), Loader=yaml.FullLoader)\n",
    "class_names = result[\"names\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "person [743  47 397 670]\n",
      "tie [442 437  54 272]\n",
      "person [125 192 585 520]\n",
      "tie [978 313  46 101]\n",
      "Detections: 4\n"
     ]
    }
   ],
   "source": [
    "# 调色板\n",
    "colors = [(255, 255, 0), (0, 255, 0), (0, 255, 255), (255, 0, 0)]\n",
    "# 显示检测框bbox\n",
    "\n",
    "for (class_id, confidence, box) in zip(filtered_ids, filered_confidences, filtered_boxes):\n",
    "    color = colors[int(class_id) % len(colors)]\n",
    "    cv2.rectangle(img, box, color, 2)\n",
    "    cv2.rectangle(img, (box[0], box[1] - 20), (box[0] + box[2], box[1]), color, -1)\n",
    "    cv2.putText(img, class_names[class_id], (box[0], box[1] - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)\n",
    "    print(class_names[class_id],box)\n",
    "    \n",
    "print(\"Detections: \" + str(len(filtered_ids)))\n",
    "cv2.imshow(\"YOLOv5+OpenVINO Demo\", img)\n",
    "cv2.waitKey(0)\n",
    "cv2.destroyAllWindows()"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "6933b49e6757033527a4c2e0a0b389a38abfc44a27af3e3e69c1b8f2d07db521"
  },
  "kernelspec": {
   "display_name": "Python 3.8.12 ('ptov')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.12"
  },
  "orig_nbformat": 2
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
